#include "node_struct.h"

__device__ float P(unsigned char i, unsigned char j, float l){
	return i == j ? 0.25f + 0.75f * __expf(-4.0f * 0.1f * l) : 0.25f - 0.25f * __expf(-4.0f * 0.1f * l);
}

__global__ void UpdateLengthDeltas_Kernel(pNode* pNodeArray, int treeLength, int* globalBest, int* neighbourhoodBest, float4 swarmConstants, float3 randomVectors){
	const int p = gridDim.x * blockIdx.x + threadIdx.x;
	const int swarmSize = gridDim.x * blockDim.x;
	//update deltas
	for(int l = 0; l < treeLength; l++){
		pNodeArray[l * swarmSize + p].delta = swarmConstants.w * pNodeArray[l * swarmSize + p].delta;										//local component
		pNodeArray[l * swarmSize + p].delta = swarmConstants.x * randomVectors.x * pNodeArray[l * swarmSize + p].lengthBest;				//local best
		pNodeArray[l * swarmSize + p].delta = swarmConstants.y * randomVectors.y * pNodeArray[l * swarmSize + neighbourhoodBest[p]].delta;	//neighbourhood best
		pNodeArray[l * swarmSize + p].delta = swarmConstants.z * randomVectors.z * pNodeArray[l * swarmSize + *globalBest].delta;			//global best
	}
	__syncthreads();
	//update lengths
	for(int l = 0; l < treeLength; l++)
		pNodeArray[l * swarmSize + p].length += pNodeArray[l * swarmSize + p].delta;
}

__global__ void NeighbourMax_Kernel(float* likelihood, float* likelihoodBest, int* neighbourhood, int neighbourhoodSize, int* neighbourhoodBest, pNode* pNodeArray, int treeLength){
	const int p = gridDim.x * blockIdx.x + threadIdx.x;
	const int swarmSize = gridDim.x * blockDim.x;
	//update local best
	if(likelihood[p] > likelihoodBest[p]){
		for(int l = 0; l < treeLength; l++)
			pNodeArray[l * swarmSize + p].lengthBest = pNodeArray[l * swarmSize + p].length;
		likelihoodBest[p] = likelihood[p];
		likelihood[p] = 0.0f;
	}
	__syncthreads();
	//update neighbour best
	for(int n = 0; n < neighbourhoodSize; n++)
		if(neighbourhood[n * swarmSize + p] > neighbourhoodBest[p])
			neighbourhoodBest[p] = neighbourhood[n * swarmSize + p];
}

__global__ void Maximum_Kernel(float* likelihood, int* globalMax){
	const int swarmSize = gridDim.x * blockDim.x;
	*globalMax = 0;
	//iterate over each particle in swarm
	for(int p = 0; p < swarmSize; p++)
		if(likelihood[p] > likelihood[*globalMax])
			*globalMax = p;
}

__global__ void EvaluateParticles_Kernel(cNode* nodeArray,
										 int treeLength,
										 pNode* pNodeArray,
										 float priorProbability[4],
										 float* likelihood
										 ){
	//indexer for current node in recursion - try and keep in local memory or promote to shared
	const unsigned int pcount = gridDim.x * blockDim.x;
	unsigned int particleIdx = blockDim.x * blockIdx.x + threadIdx.x;
	unsigned int nodeIdx;
	unsigned int particleLeftIdx;
	unsigned int particleRightIdx;

	//tree is kept in shared memory for access by all threads in block
	__shared__ unsigned char idx;
	__shared__ bool lp;
	__shared__ bool brk;
	
	if(threadIdx.x == 0){
		idx = 0;
		brk = false;
		lp = true;
	}
	__syncthreads();

	extern __shared__ cNode _nodeArray[];

	//copy tree and state to shared memory
	if(threadIdx.x == 0)
		for(int i = 0; i < treeLength; i++)
			_nodeArray[i] = nodeArray[i];
	__syncthreads();

	//each thread copies the tree tip state to their own state array
	for(int i = 0; i < treeLength; i++){
		nodeIdx = i * pcount + particleIdx;
		pNodeArray[nodeIdx].likelihood = 0.0f;
		pNodeArray[nodeIdx].state[0] = _nodeArray[i].state == 0;
		pNodeArray[nodeIdx].state[1] = _nodeArray[i].state == 1;
		pNodeArray[nodeIdx].state[2] = _nodeArray[i].state == 2;
		pNodeArray[nodeIdx].state[3] = _nodeArray[i].state == 3;
	}

	//loop fascilitates simulated recursion
	while(lp)
	{
		__syncthreads();
		if(threadIdx.x == 0) brk = false;

		//recurse left
		if(threadIdx.x == 0 && _nodeArray[idx].left != idx && _nodeArray[idx].recurse == 0){
			_nodeArray[idx].recurse = 1;
			idx = _nodeArray[idx].left;
			brk = true;
		}
		__syncthreads();
		if(brk)continue;

		//recurse right
		if(threadIdx.x == 0 && _nodeArray[idx].right != idx && _nodeArray[idx].recurse == 1){
			_nodeArray[idx].recurse = 2;
			idx = _nodeArray[idx].right;
			brk = true;
		}
		__syncthreads();
		if(brk)continue;
		
		//determine indexes
		nodeIdx = idx * pcount + particleIdx;
		particleLeftIdx = _nodeArray[idx].left * pcount + particleIdx;
		particleRightIdx = _nodeArray[idx].right * pcount + particleIdx;

		//sum across all states
		pNodeArray[nodeIdx].likelihood += P(_nodeArray[_nodeArray[idx].parent].state, _nodeArray[idx].state, pNodeArray[nodeIdx].length) * pNodeArray[nodeIdx].state[_nodeArray[idx].state];

		////determine 
		if(_nodeArray[idx].recurse == 2)
		{
			//final calculation for determining likelihood of subtree
			pNodeArray[nodeIdx].state[_nodeArray[idx].state] = pNodeArray[particleLeftIdx].likelihood * pNodeArray[particleRightIdx].likelihood;
			//reset the values of the children
			pNodeArray[particleLeftIdx].likelihood = pNodeArray[particleRightIdx].likelihood = 0.0f;
		}
		
		__syncthreads();

		//reset recursion flags so the tree can be reused on next traversal
		if(threadIdx.x == 0 && _nodeArray[idx].recurse == 2)
			_nodeArray[idx].recurse = 0;
		__syncthreads();


		//go to next state
		if(threadIdx.x == 0 && _nodeArray[idx].state != 3){
			_nodeArray[idx].state++;
			brk = true;
		}
		__syncthreads();
		if(brk) continue;

		//return to parent
		if(threadIdx.x == 0 && _nodeArray[idx].state == 3 && _nodeArray[idx].parent != idx){
			_nodeArray[idx].state = 0;
			idx = nodeArray[idx].parent;
			brk = true;
		}
		__syncthreads();
		if(brk) continue;

		if(threadIdx.x == 0) lp = false;
		__syncthreads();
	}
	
	//Summate over all states
	likelihood[particleIdx] += pNodeArray[particleIdx].state[0] * priorProbability[0];
	likelihood[particleIdx] += pNodeArray[particleIdx].state[1] * priorProbability[1];
	likelihood[particleIdx] += pNodeArray[particleIdx].state[2] * priorProbability[2];
	likelihood[particleIdx] += pNodeArray[particleIdx].state[3] * priorProbability[3];
}

